#pragma once

#define GLM_FORCE_RADIANS
#include "glm/glm.hpp"
#include "glm/gtc/constants.hpp"
#include "glm/gtc/matrix_transform.hpp"

#define VK_USE_PLATFORM_WIN32_KHR
#include <vulkan/vulkan.h>

#include <vector>
#include <string>
#include "../Utils/PlatTypeUtil.h"
#include "../Templates/RefCounting.h"
#include <unordered_map>
#include <array>


class RenderState;

class VulkanRHI
{
public:

	VulkanRHI();
	~VulkanRHI();

	void Init();

	void Exit();

	VkInstance GetInstance() const;

	VkPhysicalDevice GetPhysicalDevice() const;

	VkDevice GetLogicalDevice() const;

	VkSurfaceKHR GetSurface() const;

	VkFormat GetDepthFormat() const;

	VkSampleCountFlagBits GetMsaaSamples() const;

	VkPhysicalDeviceProperties GetDeviceProperties() const;

	VkPhysicalDeviceFeatures GetDeviceFeatures() const;

	VkQueue GetQueue(VkQueueFlagBits queueFlags = VK_QUEUE_GRAPHICS_BIT) const;

	VkQueue GetPresentQueue() const;

	uint32 GetQueueFamily(VkQueueFlagBits queueFlags = VK_QUEUE_GRAPHICS_BIT) const;

	uint32 GetPresentFamily() const;

	VkPhysicalDeviceRayTracingPipelinePropertiesKHR GetRayTracingPipelineProperties() const;

	VkPhysicalDeviceAccelerationStructureFeaturesKHR GetAccelerationStructureFeatures() const;

	uint32 FindMemoryType(uint32 typeFilter, VkMemoryPropertyFlags properties);

public:

	void GuessAccessByLayout(
		VkAccessFlags& srcAccessMask,
		VkAccessFlags& dstAccessMask,
		VkImageLayout oldLayout,
		VkImageLayout newLayout);

	void CreateBuffer(
		VkBuffer& buffer,
		VkDeviceMemory& bufferMemory,
		VkDeviceSize size,
		VkBufferUsageFlags usage,
		VkMemoryPropertyFlags memoryProperty,
		VkSharingMode sharingMode = VK_SHARING_MODE_EXCLUSIVE,
		std::vector<VkQueueFlagBits> queueFamilys = std::vector<VkQueueFlagBits>());

	void CreateImage(
		VkImage& image,
		VkDeviceMemory& imageMemory,
		uint32 width,
		uint32 height,
		VkImageUsageFlags usage,
		VkFormat format = VK_FORMAT_R8G8B8A8_UNORM,
		uint32 mipLevels = 1,
		VkMemoryPropertyFlags memoryProperty = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
		VkSampleCountFlagBits samples = VK_SAMPLE_COUNT_1_BIT,
		VkImageType imageType = VK_IMAGE_TYPE_2D,
		VkImageLayout imageLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		uint32 arrayLayers = 1,
		VkImageTiling tiling = VK_IMAGE_TILING_OPTIMAL,
		VkSharingMode sharingMode = VK_SHARING_MODE_EXCLUSIVE,
		std::vector<VkQueueFlagBits> queueFamilys = std::vector<VkQueueFlagBits>());

	void CreateImage(
		VkImage& image,
		VkDeviceMemory& imageMemory,
		VkImageCreateFlags flags,
		VkExtent3D extent,
		VkImageUsageFlags usage,
		VkFormat format = VK_FORMAT_R8G8B8A8_UNORM,
		uint32 mipLevels = 1,
		VkMemoryPropertyFlags memoryProperty = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
		VkSampleCountFlagBits samples = VK_SAMPLE_COUNT_1_BIT,
		VkImageType imageType = VK_IMAGE_TYPE_2D,
		VkImageLayout imageLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		uint32 arrayLayers = 1,
		VkImageTiling tiling = VK_IMAGE_TILING_OPTIMAL,
		VkSharingMode sharingMode = VK_SHARING_MODE_EXCLUSIVE,
		std::vector<VkQueueFlagBits> queueFamilys = std::vector<VkQueueFlagBits>());

	void CreateImage(
		VkImage& image,
		VkDeviceMemory& imageMemory,
		VkExtent3D extent,
		VkImageUsageFlags usage,
		VkFormat format = VK_FORMAT_R8G8B8A8_UNORM,
		uint32 mipLevels = 1,
		VkMemoryPropertyFlags memoryProperty = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
		VkSampleCountFlagBits samples = VK_SAMPLE_COUNT_1_BIT,
		VkImageType imageType = VK_IMAGE_TYPE_2D,
		VkImageLayout imageLayout = VK_IMAGE_LAYOUT_UNDEFINED,
		uint32 arrayLayers = 1,
		VkImageTiling tiling = VK_IMAGE_TILING_OPTIMAL,
		VkSharingMode sharingMode = VK_SHARING_MODE_EXCLUSIVE,
		std::vector<VkQueueFlagBits> queueFamilys = std::vector<VkQueueFlagBits>());

	void CreateImageView(
		VkImageView& imageView,
		VkImage image,
		VkFormat format = VK_FORMAT_R8G8B8A8_UNORM,
		VkImageViewType viewType = VK_IMAGE_VIEW_TYPE_2D,
		VkImageSubresourceRange subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT , 0, VK_REMAINING_MIP_LEVELS, 0, VK_REMAINING_ARRAY_LAYERS },
		VkComponentMapping components = { VK_COMPONENT_SWIZZLE_R , VK_COMPONENT_SWIZZLE_G , VK_COMPONENT_SWIZZLE_B , VK_COMPONENT_SWIZZLE_A });

	void TransferLayout(
		VkCommandBuffer commandBuffer,
		VkImage image,
		VkImageLayout oldLayout,
		VkImageLayout newLayout,
		VkImageSubresourceRange subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT , 0, VK_REMAINING_MIP_LEVELS, 0, VK_REMAINING_ARRAY_LAYERS },
		VkPipelineStageFlags srcStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT,
		VkPipelineStageFlags dstStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT,
		VkAccessFlags srcAccessMask = VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM,
		VkAccessFlags dstAccessMask = VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM,
		VkQueueFlagBits srcQueueFlag = VK_QUEUE_FLAG_BITS_MAX_ENUM,
		VkQueueFlagBits dstQueueFlag = VK_QUEUE_FLAG_BITS_MAX_ENUM);

	void TransferLayout(
		VkImage image,
		VkImageLayout oldLayout,
		VkImageLayout newLayout,
		VkImageSubresourceRange subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT , 0, VK_REMAINING_MIP_LEVELS, 0, VK_REMAINING_ARRAY_LAYERS },
		VkQueueFlagBits commandQueue = VK_QUEUE_GRAPHICS_BIT,
		VkPipelineStageFlags srcStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT,
		VkPipelineStageFlags dstStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT,
		VkAccessFlags srcAccessMask = VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM,
		VkAccessFlags dstAccessMask = VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM,
		VkQueueFlagBits srcQueueFlag = VK_QUEUE_FLAG_BITS_MAX_ENUM,
		VkQueueFlagBits dstQueueFlag = VK_QUEUE_FLAG_BITS_MAX_ENUM);

	void CopyBufferToImage(
		VkCommandBuffer commandBuffer,
		VkBuffer buffer,
		VkImage image,
		uint32 width,
		uint32 height,
		VkImageSubresourceLayers imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, VK_REMAINING_ARRAY_LAYERS },
		VkDeviceSize bufferOffset = 0,
		uint32 bufferRowLength = 0,
		uint32 bufferImageHeight = 0,
		VkOffset3D imageOffset = { 0, 0, 0 },
		VkImageLayout dstImageLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

	void CopyBufferToImage(
		VkCommandBuffer commandBuffer, 
		VkBuffer buffer,
		VkImage image,
		VkExtent3D extent,
		VkImageSubresourceLayers imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, VK_REMAINING_ARRAY_LAYERS },
		VkDeviceSize bufferOffset = 0,
		uint32 bufferRowLength = 0,
		uint32 bufferImageHeight = 0,
		VkOffset3D imageOffset = { 0, 0, 0 },
		VkImageLayout dstImageLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

	void CopyBufferToImage(
		VkBuffer buffer,
		VkImage image,
		uint32 width,
		uint32 height,
		VkImageSubresourceLayers imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, VK_REMAINING_ARRAY_LAYERS },
		VkQueueFlagBits commandQueue = VK_QUEUE_GRAPHICS_BIT,
		VkDeviceSize bufferOffset = 0,
		uint32 bufferRowLength = 0,
		uint32 bufferImageHeight = 0,
		VkOffset3D imageOffset = { 0, 0, 0 },
		VkImageLayout dstImageLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

	void CopyBufferToImage(
		VkBuffer buffer,
		VkImage image,
		VkExtent3D extent,
		VkImageSubresourceLayers imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, VK_REMAINING_ARRAY_LAYERS },
		VkQueueFlagBits commandQueue = VK_QUEUE_GRAPHICS_BIT,
		VkDeviceSize bufferOffset = 0,
		uint32 bufferRowLength = 0,
		uint32 bufferImageHeight = 0,
		VkOffset3D imageOffset = { 0, 0, 0 },
		VkImageLayout dstImageLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

	void CopyImageToBuffer(
		VkCommandBuffer commandBuffer,
		VkImage image,
		VkBuffer buffer,
		VkExtent3D extent,
		VkImageSubresourceLayers imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, VK_REMAINING_ARRAY_LAYERS },
		VkDeviceSize bufferOffset = 0,
		uint32 bufferRowLength = 0,
		uint32 bufferImageHeight = 0,
		VkOffset3D imageOffset = { 0, 0, 0 },
		VkImageLayout srcImageLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
		);

	void CopyImageToBuffer(
		VkImage image,
		VkBuffer buffer,
		VkExtent3D extent,
		VkImageSubresourceLayers imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, VK_REMAINING_ARRAY_LAYERS },
		VkQueueFlagBits commandQueue = VK_QUEUE_GRAPHICS_BIT,
		VkDeviceSize bufferOffset = 0,
		uint32 bufferRowLength = 0,
		uint32 bufferImageHeight = 0,
		VkOffset3D imageOffset = { 0, 0, 0 },
		VkImageLayout srcImageLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
	);
	
	// Generate mipmap for texture2D, but the layers of mip = 0 should have data
	void GenerateMipmaps2D(
		VkCommandBuffer commandBuffer,
		VkImage image,
		uint32 width,
		uint32 height,
		VkFormat format,
		uint32 mipLevels,
		uint32 arrayLayers,
		VkImageLayout oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		VkImageLayout newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
		VkAccessFlags dstAccessMask = VK_ACCESS_SHADER_READ_BIT,
		VkPipelineStageFlags dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
		VkImageAspectFlags aspectMask = VK_IMAGE_ASPECT_COLOR_BIT);

	void GenerateMipmaps2D(
		VkImage image,
		uint32 width,
		uint32 height,
		VkFormat format,
		uint32 mipLevels,
		uint32 arrayLayers,
		VkQueueFlagBits commandQueue = VK_QUEUE_GRAPHICS_BIT,
		VkImageLayout oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		VkImageLayout newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
		VkAccessFlags dstAccessMask = VK_ACCESS_SHADER_READ_BIT,
		VkPipelineStageFlags dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
		VkImageAspectFlags aspectMask = VK_IMAGE_ASPECT_COLOR_BIT);


protected:

	void CreateInstance();

	void SetupReportCallback();

	void DestroyReportCallback();

	void CreateSurface();

	void SelectPhysicalDevice();

	void CreateLogicalDevice();

	void CreateUniversal();

protected:

	uint32 mApiVersion;
	bool mEnableValidation;
	VkInstance mInstance;
	std::vector<std::string> mSupportedInstanceExtensions;
	std::vector<const char*> mEnabledInstanceExtensions;
	std::vector<const char*> mValidationLayers;
	std::vector<const char*> mDeviceExtensions;

	VkDebugUtilsMessengerEXT mDebugMessenger;

	VkSurfaceKHR mSurface;

	VkPhysicalDevice mPhysicalDevice;

	VkPhysicalDeviceProperties2 mDeviceProperties2;
	VkPhysicalDeviceFeatures2 mDeviceFeatures2;

	VkPhysicalDeviceProperties mDeviceProperties10;
	VkPhysicalDeviceFeatures mDeviceFeatures10;

	VkPhysicalDeviceVulkan11Properties mDeviceProperties11;
	VkPhysicalDeviceVulkan11Features mDeviceFeatures11;

	VkPhysicalDeviceVulkan12Properties mDeviceProperties12;
	VkPhysicalDeviceVulkan12Features mDeviceFeatures12;

	// Addition properties and features
	VkPhysicalDeviceRayTracingPipelinePropertiesKHR  mRayTracingPipelineProperties;

	VkPhysicalDeviceAccelerationStructureFeaturesKHR mAccelerationStructureFeatures;
	VkPhysicalDeviceRayTracingPipelineFeaturesKHR mRayTracingPipelineFeatures;
	VkPhysicalDeviceBufferDeviceAddressFeaturesKHR mBufferDeviceAddressFeatures;
	VkPhysicalDeviceDescriptorIndexingFeatures mDescriptorIndexingFeatures;
	VkPhysicalDeviceRayQueryFeaturesKHR mRayQueryFeatures;



	uint32 mGraphicsFamily;
	uint32 mComputeFamily;
	uint32 mTransferFamily;
	uint32 mPresentFamily;
	VkSampleCountFlagBits mMsaaSamples;

	VkDevice mLogicalDevice;

	VkQueue mGraphicsQueue;
	VkQueue mComputeQueue;
	VkQueue mTransferQueue;
	VkQueue mPresentQueue;

	VkPhysicalDeviceMemoryProperties mMemoryProperties;

	VkFormat mDepthFormat;
};